﻿#include "arp.h"
#include "mblock.h"
#include "protocol.h"
#include "tools.h"
#include "ether.h"
#include "timer.h"
#include "ipv4.h"

static arp_entry_t cache_tbl[ARP_CACHE_SIZE]; // 用于存储ARP表项的空间
static mblock_t cache_block;                  // 内存管理器
static nlist_t cache_list;
static const uint8_t empty_hwaddr[] = {0, 0, 0, 0, 0, 0};
static net_timer_t cache_timer; // arp缓存表项定时器,用于定期刷新arp缓存表项

#if DBG_DISP_ENABLED(DBG_ARP_LEVEL)
/**
 * @brief 打印arp包的内容
 */
static void arp_pkt_display(arp_pkt_t *arp_pkt)
{
    plat_printf("-----------------arp packet--------------------\n");
    plat_printf("htype: %d, hlen = %d\n", x_ntohs(arp_pkt->htype), arp_pkt->hwlen);
    plat_printf("ptype: %0x, plen = %d\n", x_ntohs(arp_pkt->ptype), arp_pkt->plen);
    plat_printf("opcode: %d", x_ntohs(arp_pkt->opcode));
    switch (x_ntohs(arp_pkt->opcode))
    {
    case ARP_REQUEST:
        plat_printf("  type: request\n");
        break;
    case ARP_REPLY:
        plat_printf("  type: reply\n");
        break;
    default:
        plat_printf("  type: unknown\n");
        break;
    }

    ipaddr_t ipaddr = ipaddr_from_buf(arp_pkt->sender_ipaddr, IPADDR_V4);
    dbg_dump_ipaddr("sender ip: ", &ipaddr);
    dbg_dump_hwaddr("sender hwaddr: ", arp_pkt->sender_hwaddr, ETHER_HWA_SIZE);
    ipaddr = ipaddr_from_buf(arp_pkt->target_ipaddr, IPADDR_V4);
    dbg_dump_ipaddr("target ip: ", &ipaddr);
    dbg_dump_hwaddr("target hwaddr: ", arp_pkt->target_hwaddr, ETHER_HWA_SIZE);
    plat_printf("----------------------------------------------\n");
}

/**
 * @brief 打印arp表项
 */
static void display_arp_entry(const arp_entry_t *entry)
{
    plat_printf("*************index: %d*************\n", (int)(entry - cache_tbl));
    ipaddr_t ipaddr = ipaddr_from_buf(entry->ipaddr, IPADDR_V4);
    dbg_dump_ipaddr("ipaddr: ", &ipaddr);
    dbg_dump_hwaddr("hwaddr: ", entry->hwaddr, ETHER_HWA_SIZE);
    plat_printf("tmo: %d, retry: %d, state: %d, buf_size: %d\n", entry->tmo, entry->retry, entry->state, nlist_count(&entry->buf_list));
    plat_printf("***********************************\n");
}

/**
 * @brief arp缓存表打印
 */
static void display_arp_tbl()
{
    // arp表是动态的,所以直接遍历存放arp的数组
    plat_printf("---------------------arp tbl----------------------\n");
    arp_entry_t *entry = cache_tbl;
    for (int i = 0; i < ARP_CACHE_SIZE; ++i)
    {
        if (entry->state == NET_ARP_RESOLVED || entry->state == NET_ARP_WAITING)
        {
            display_arp_entry(entry++);
        }
    }
    plat_printf("--------------------------------------------------\n");
}
#else
#define arp_pkt_display(arp_pkt)
#define display_arp_tbl()
#define display_arp_entry(entry)
#endif

/**
 * @brief 将以ms为单位的时间转化为扫描次数
 */
static int arp_tmo_to_cnt(int tmo)
{
    return tmo / ARP_TIMER_TMO;
}

/**
 * @移除指定arp表项中缓存链表中的包
 */
static void cache_clear_all(arp_entry_t *entry)
{
    nlist_node_t *first = NULL;
    while (first = nlist_remove_first(&entry->buf_list))
    {
        pktbuf_t *buf = nlist_entry(pktbuf_t, node, first);
        pktbuf_free(buf);
    }
}

/**
 * @brief 删除指定arp缓存表项
 */
static void cache_free(arp_entry_t *entry)
{
    cache_clear_all(entry);
    nlist_remove(&cache_list, &entry->node);
    mblock_free(&cache_block, entry);
}

/**
 * @brief 申请arp表项
 * @param force 是否在内存表项空间不足时强制分配
 * @return 申请的arp表项内存
 */
static arp_entry_t *cache_alloc(bool force)
{
    arp_entry_t *entry = mblock_alloc(&cache_block, -1);
    // 分配空间失败时,考虑是否强制分配
    if (!entry && force)
    {
        DBG_INFO(DBG_ARP_LEVEL, "alloc arp entry failed and will insert forcely!");
        nlist_node_t *node = nlist_remove_last(&cache_list);
        if (!node)
        {
            DBG_WARNING(DBG_ARP_LEVEL, "remove last node failed in cache list! list size: %d!", nlist_count(&cache_list));
            return NULL;
        }
        entry = nlist_entry(arp_entry_t, node, node);
        // 释放表项之前累积未处理的包 todo...
    }
    else
    {
        // 直接分配成功
        plat_memset(entry, 0, sizeof(entry));
        entry->state = NET_ARP_FREE;
        nlist_node_init(&entry->node);
        nlist_init(&entry->buf_list);
    }

    return entry;
}

/**
 * @brief 根据ip地址查找对应的arp表项
 * @param ip ip地址
 * @return 查找到的arp表项,未找到则为NULL
 */
static arp_entry_t *cache_find(const uint8_t *ip)
{
    nlist_node_t *p;
    nlist_for_each(p, &cache_list)
    {
        arp_entry_t *curr_entry = nlist_entry(arp_entry_t, node, p);
        if (plat_memcmp(curr_entry->ipaddr, ip, IPV4_ADDR_SIZE) == 0)
        {
            // 查到到后放置在表的前面,以提高速率
            nlist_remove(&cache_list, p);
            nlist_insert_first(&cache_list, p);
            return curr_entry;
        }
    }

    return NULL;
}

/**
 * @brief 对表项的各个参数进行设置
 */
static void cache_entry_set(arp_entry_t *entry, const uint8_t *hwaddr, const uint8_t *ip, netif_t *netif, int state)
{
    plat_memcpy(entry->hwaddr, hwaddr, ETHER_HWA_SIZE);
    plat_memcpy(entry->ipaddr, ip, IPV4_ADDR_SIZE);
    entry->netif = netif;
    entry->state = state;
    if (state == NET_ARP_RESOLVED)
    {
        entry->tmo = arp_tmo_to_cnt(ARP_ENTRY_STABLE_TMO);
    }
    else
    {
        entry->tmo = arp_tmo_to_cnt(ARP_ENTRY_PENDING);
    }
    entry->retry = ARP_ENTRY_RETRY_CNT;
}

/**
 * @brief 将表项中的缓存包全部发送出去
 */
static net_err_t cache_send_all(arp_entry_t *entry)
{
    ipaddr_t ipaddr;
    ipaddr = ipaddr_from_buf(entry->ipaddr, IPADDR_V4);
    // dbg_dump_ipaddr("send all pkt in arp_entry of netif:%s to ip: ", &ipaddr);

    nlist_node_t *first;
    while (first = nlist_remove_first(&entry->buf_list))
    {
        pktbuf_t *buf = nlist_entry(pktbuf_t, node, first);
        net_err_t err = ether_raw_out(entry->netif, NET_PROTOCOL_IPV4, entry->hwaddr, buf);
        if (err < NET_ERR_OK)
        {
            DBG_INFO(DBG_ARP_LEVEL, "send pkt in entry list by netif:%s failed!", entry->netif->name);
            pktbuf_free(buf);
        }
    }

    return NET_ERR_OK;
}

/**
 * @brief 向arp表中插入表项
 * @param netif 需要该表项的网络接口
 * @param ip ip地址
 * @param hwaddr 硬件地址
 * @param force 内存空间不足时，是否强制插入
 */
static net_err_t cache_insert(netif_t *netif, const uint8_t *ip, const uint8_t *hwaddr, bool force)
{
    if (*(uint32_t *)ip == 0)
    {
        DBG_WARNING(DBG_ARP_LEVEL, "the arp ipaddr is null!");
        return NET_ERR_NOSUPPORT;
    }

    // 先查找该表项存不存在
    arp_entry_t *entry = cache_find(ip);
    if (!entry)
    {
        // 表项不存在则申请空间并插入
        DBG_INFO(DBG_ARP_LEVEL, "insert a new arp entry!");
        entry = cache_alloc(force);
        if (!entry)
        {
            DBG_WARNING(DBG_ARP_LEVEL, "alloc entry to ip failed!");
            return NET_ERR_NULLP;
        }
        cache_entry_set(entry, hwaddr, ip, netif, NET_ARP_RESOLVED);
        nlist_insert_first(&cache_list, &entry->node);
    }
    else
    {
        // 表项已经存在则进行更新
        DBG_INFO(DBG_ARP_LEVEL, "arp entry is existed, update it!");
        cache_entry_set(entry, hwaddr, ip, netif, NET_ARP_RESOLVED);
        if (nlist_first(&cache_list) != &entry->node)
        {
            nlist_remove(&cache_list, &entry->node);
            nlist_insert_first(&cache_list, &entry->node);
        }
        cache_send_all(entry);
    }

    display_arp_tbl();
    return NET_ERR_OK;
}

/**
 * @brief 缓存表相关初始化
 */
static net_err_t cache_init()
{
    nlist_init(&cache_list);
    plat_memset(cache_tbl, 0, sizeof(arp_entry_t) * ARP_CACHE_SIZE);
    net_err_t err = mblock_init(&cache_block, cache_tbl, sizeof(arp_entry_t), ARP_CACHE_SIZE, NLOCKER_NONE);
    if (err < NET_ERR_OK)
    {
        DBG_ERROR(DBG_ARP_LEVEL, "mblock init failed in arp!");
    }
    return err;
}

/**
 * @brief 周期性的对arp缓存表进行扫描
 */
static void arp_cache_tmo(net_timer_t *timer, void *arg)
{
    DBG_INFO(DBG_ARP_LEVEL, "=======refresh arp tbl!========");

    nlist_node_t *p;
    nlist_for_each(p, &cache_list)
    {
        arp_entry_t *entry = nlist_entry(arp_entry_t, node, p);
        if (--entry->tmo > 0)
            continue; // 没超时就不用处理

        // 根据缓存表项不同的状态进行不同的处理
        switch (entry->state)
        {
        case NET_ARP_RESOLVED:
        {
            // 处于有效状态,则需要进行更新
            DBG_INFO(DBG_ARP_LEVEL, "arp entry change state to pending!");
            display_arp_entry(entry);

            ipaddr_t target_ipaddr = ipaddr_from_buf(entry->ipaddr, IPADDR_V4);
            arp_make_request(entry->netif, &target_ipaddr);
            entry->tmo = arp_tmo_to_cnt(ARP_ENTRY_PENDING);
            entry->retry = ARP_ENTRY_RETRY_CNT;
            break;
        }
        case NET_ARP_WAITING:
        {
            // 处于等待，则进行重发
            DBG_INFO(DBG_ARP_LEVEL, "arp entry is watting, resend arp request!");
            if (--entry->retry == 0)
            {
                // 重发次数用完
                DBG_INFO(DBG_ARP_LEVEL, "arp entry pending timeout, free it!");
                display_arp_entry(entry);
                cache_free(entry);
            }
            else
            {
                // 还可以继续重发
                DBG_INFO(DBG_ARP_LEVEL, "send request arp again!");
                entry->tmo = arp_tmo_to_cnt(ARP_ENTRY_PENDING);
                ipaddr_t target_ipaddr = ipaddr_from_buf(entry->ipaddr, IPADDR_V4);
                arp_make_request(entry->netif, &target_ipaddr);
            }
            break;
        }
        default:
            DBG_ERROR(DBG_ARP_LEVEL, "the arp entry is error!");
            display_arp_entry(entry);
            break;
        }
    }
}

/**
 * @brief ARP初始化相关
 */
net_err_t arp_init()
{
    DBG_INFO(DBG_ARP_LEVEL, "arp init begin!");

    net_err_t err = cache_init();
    if (err < NET_ERR_OK)
    {
        DBG_ERROR(DBG_ARP_LEVEL, "cache init failed!");
        return err;
    }
    err = timer_add(&cache_timer, "arp_timer", arp_cache_tmo, NULL, ARP_TIMER_TMO, true);
    if (err < NET_ERR_OK)
    {
        DBG_ERROR(DBG_ARP_LEVEL, "arp add timer failed!");
        return err;
    }

    DBG_INFO(DBG_ARP_LEVEL, "arp init end!");
    return NET_ERR_OK;
}

/**
 * @brief 发送arp请求
 * @param netif 用来发送请求的网络接口
 * @param ipaddr 目的ip地址
 */
net_err_t arp_make_request(netif_t *netif, const ipaddr_t *ipaddr)
{
    pktbuf_t *buf = pktbuf_alloc(sizeof(arp_pkt_t));
    if (!buf)
    {
        DBG_ERROR(DBG_ARP_LEVEL, "alloc pkt failed in arp make request!");
        return NET_ERR_NULLP;
    }

    pktbuf_set_cont(buf, sizeof(arp_pkt_t));
    arp_pkt_t *arp_pkt = (arp_pkt_t *)pktbuf_data(buf);
    arp_pkt->htype = x_htons(ARP_HW_ETHER);
    arp_pkt->ptype = x_htons(NET_PROTOCOL_IPV4);
    arp_pkt->hwlen = ETHER_HWA_SIZE;
    arp_pkt->plen = IPV4_ADDR_SIZE;
    arp_pkt->opcode = x_htons(ARP_REQUEST);
    plat_memcpy(arp_pkt->sender_hwaddr, netif->hwaddr.addr, ETHER_HWA_SIZE);
    ipaddr_to_buf(&netif->ipaddr, arp_pkt->sender_ipaddr);
    plat_memcpy(arp_pkt->target_hwaddr, ether_broadcast_addr(), ETHER_HWA_SIZE);
    ipaddr_to_buf(ipaddr, arp_pkt->target_ipaddr);

    net_err_t err = ether_raw_out(netif, NET_PROTOCOL_ARP, arp_pkt->target_hwaddr, buf);
    if (err < 0)
    {
        pktbuf_free(buf);
        DBG_ERROR(DBG_ARP_LEVEL, "send arp request by netif:%s failed!", netif->name);
        return NET_ERR_SYS;
    }

    return NET_ERR_OK;
}

/**
 * @brief 回应arp请求
 * @param netif 回应arp请求的网络接口
 * @param buf 收到的arp请求包
 */
net_err_t arp_make_reply(netif_t *netif, pktbuf_t *buf)
{
    arp_pkt_t *pkt = (arp_pkt_t *)pktbuf_data(buf);
    // 直接将收到的包改变操作类型 和 收发地址返回即可
    pkt->opcode = x_htons(ARP_REPLY);
    plat_memcpy(pkt->target_hwaddr, pkt->sender_hwaddr, ETHER_HWA_SIZE);
    plat_memcpy(pkt->target_ipaddr, pkt->sender_ipaddr, IPV4_ADDR_SIZE);
    plat_memcpy(pkt->sender_hwaddr, netif->hwaddr.addr, ETHER_HWA_SIZE);
    ipaddr_to_buf(&netif->ipaddr, pkt->sender_ipaddr);

    net_err_t err = ether_raw_out(netif, NET_PROTOCOL_ARP, pkt->target_hwaddr, buf);
    if (err < NET_ERR_OK)
    {
        DBG_INFO(DBG_ARP_LEVEL, "reply arp failed!");
        pktbuf_free(buf);
    }

    return err;
}

/**
 * @brief 向外发送无回报arp包
 * @param netif 用于发送arp包的网络接口
 */
net_err_t arp_make_gratuitous(netif_t *netif)
{
    return arp_make_request(netif, &netif->ipaddr);
}

/**
 * @brief 判断收到的arp包是否有效
 * @param pkt 收到的arp包
 * @param size 包的大小
 * @param netif 收到包的网络接口
 */
static bool arp_is_pkt_ok(arp_pkt_t *pkt, uint16_t size, netif_t *netif)
{
    // 判断包的大小是否合规
    if (size < sizeof(arp_pkt_t))
        return false;

    // 判断包的参数能否支持
    if ((x_ntohs(pkt->htype) != ARP_HW_ETHER) ||
        (pkt->hwlen != ETHER_HWA_SIZE) ||
        (x_ntohs(pkt->ptype) != NET_PROTOCOL_IPV4) ||
        (pkt->plen != IPV4_ADDR_SIZE))
    {

        DBG_WARNING(DBG_ARP_LEVEL, "the pkt is not support!");
        return false;
    }

    uint16_t opcode = x_ntohs(pkt->opcode);
    if (opcode != ARP_REPLY && opcode != ARP_REQUEST)
    {
        DBG_WARNING(DBG_ARP_LEVEL, "the pkt is not support!");
        return false;
    }

    return true;
}

/**
 * @brief 处理收到的arp包
 * @param netif 收到arp包的网络接口
 * @param buf 收到的arp包
 */
net_err_t arp_in(netif_t *netif, pktbuf_t *buf)
{
    DBG_INFO(DBG_ARP_LEVEL, "arp in");

    net_err_t err = pktbuf_set_cont(buf, sizeof(arp_pkt_t));
    if (err < NET_ERR_OK)
    {
        DBG_ERROR(DBG_ARP_LEVEL, "set cont of the buf size:%d failed!", buf->total_size);
        return err;
    }

    // 判断收到的包是否合法
    arp_pkt_t *arp_pkt = (arp_pkt_t *)pktbuf_data(buf);

    arp_pkt_display(arp_pkt);

    if (!arp_is_pkt_ok(arp_pkt, buf->total_size, netif))
    {
        DBG_WARNING(DBG_ARP_LEVEL, "the arp pkt is error!");
        return NET_ERR_NOSUPPORT;
    }

    // 根据收到的arp包处理相关内容
    ipaddr_t target_ipaddr = ipaddr_from_buf(arp_pkt->target_ipaddr, IPADDR_V4);
    if (ipaddr_is_equal(&target_ipaddr, &netif->ipaddr))
    {
        DBG_INFO(DBG_ARP_LEVEL, "arp request for me!");
        cache_insert(netif, arp_pkt->sender_ipaddr, arp_pkt->sender_hwaddr, true);
        if (x_ntohs(arp_pkt->opcode) == ARP_REQUEST)
        {
            return arp_make_reply(netif, buf);
        }
    }
    else
    {
        DBG_INFO(DBG_ARP_LEVEL, "recv arp pkt not for me!");
        cache_insert(netif, arp_pkt->sender_ipaddr, arp_pkt->sender_hwaddr, false);
    }
    pktbuf_free(buf);
    return NET_ERR_OK;
}

/**
 * @brief 根据指定的ip地址查找arp表
 */
arp_entry_t *arp_find(uint8_t *ip)
{
    return cache_find(ip);
}

/**
 * @brief 解析ip地址并进行相关的操作
 * @param netif 网络接口
 * @param dest 目的ip地址
 * @param buf 需要发送的包
 */
net_err_t arp_resolve(netif_t *netif, ipaddr_t *dest, pktbuf_t *buf)
{
    if (ipaddr_is_direct_broadcast(dest, &netif->netmask) || ipaddr_is_local_broadcast(dest))
    {
        return ether_raw_out(netif, NET_PROTOCOL_IPV4, ether_broadcast_addr(), buf);
    }

    // 发出前先获取目的硬件地址
    arp_entry_t *entry = arp_find(dest->addr);
    if (entry)
    {
        DBG_INFO(DBG_ETHER_LEVEL, "find a arp entry!");
        // 根据arp表的状态进行不同的操作
        if (entry->state == NET_ARP_RESOLVED)
        {
            return ether_raw_out(netif, NET_PROTOCOL_IPV4, entry->hwaddr, buf);
        }

        if (nlist_count(&entry->buf_list) < ARP_MAX_PKT_WAITING)
        {
            nlist_insert_last(&entry->buf_list, &buf->node);
        }
        else
        {
            DBG_WARNING(DBG_ARP_LEVEL, "the buf in netif:%s is too many, the list count:%d!\n", netif->name, nlist_count(&entry->buf_list));
            display_arp_tbl();
            return NET_ERR_FULL;
        }
    }
    else
    {
        // 没有找到相关的arp表项
        DBG_INFO(DBG_ARP_LEVEL, "the ipaddr can not find in arp list!");
        arp_entry_t *entry = cache_alloc(true);
        if (!entry)
        {
            DBG_ERROR(DBG_ARP_LEVEL, "alloc entry failed in arp resolve!");
            return NET_ERR_SYS;
        }

        cache_entry_set(entry, empty_hwaddr, dest->addr, netif, NET_ARP_WAITING);
        nlist_insert_last(&entry->buf_list, &buf->node);
        nlist_insert_first(&cache_list, &entry->node);

        display_arp_tbl();
        return arp_make_request(netif, dest);
    }

    return NET_ERR_OK;
}

/**
 * @brief 删除指定网络接口的arp缓存表
 * @param netif 指定的网络接口
 */
void arp_clear(netif_t *netif)
{
    nlist_node_t *p;
    nlist_for_each(p, &cache_list)
    {
        arp_entry_t *entry = nlist_entry(arp_entry_t, node, p);
        if (entry->netif == netif)
        {
            cache_free(entry);
        }
    }
}

/**
 * @brief 根据收到ip包来更新arp表
 * @param netif 收到该包的网络接口
 * @param buf 收到的包
 */
net_err_t arp_update_from_ipbuf(netif_t *netif, pktbuf_t *buf)
{
    net_err_t err = pktbuf_set_cont(buf, sizeof(ether_hdr_t) + sizeof(ipv4_hdr_t));
    if (err < NET_ERR_OK)
    {
        DBG_WARNING(DBG_ARP_LEVEL, "pktbuf set cont failed! size:%d, err_type:%d!", sizeof(ether_hdr_t) + sizeof(ipv4_hdr_t), err);
        return err;
    }

    ether_hdr_t *ether_hdr = (ether_hdr_t *)pktbuf_data(buf);
    ipv4_hdr_t *ipv4_hdr = (ipv4_hdr_t *)((uint8_t *)ether_hdr + sizeof(ether_hdr_t));

    cache_insert(netif, ipv4_hdr->src_ip, ether_hdr->src, false);
    return NET_ERR_OK;
}